from pymongo import MongoClient
from pymongo.errors import BulkWriteError
from bson.objectid import ObjectId
from datetime import datetime, timedelta
from faker import Faker
from tqdm import tqdm
import random

MONGO_URL = 'mongodb://localhost:27017/'
OID_COUNT = 5000
OPEN_ID_COUNT = 100000
DATA_SIZE = 1000000
# 连接到MongoDB数据库
mongo_client = MongoClient(MONGO_URL)


def generate_data(start_datetime, end_datetime, fake, oid_list, open_id_list):
    '''
    随机生成一条数据
    '''
    return {
        'atime': fake.date_time_between(start_date=start_datetime, end_date=end_datetime),  # 随机生成指定范围内的日期和时间
        'open_id': random.choice(open_id_list),  # 随机选择一个UUID作为用户唯一标识符
        'oid': random.choice(oid_list),  # 随机选择一个ObjectId
        'version': random.randint(1, 10),  # 随机生成一个1到10之间的APP版本号
    }



def generate_bach_data(start_datetime, end_datetime, data_size, oid_list, open_id_list, fake):
    '''
    随机生成批量数据
    '''
    data_list = []
    progress_bar = tqdm(total=data_size, desc='生成数据', unit='条')
    for _ in range(data_size):
        post_data = {
            'atime': fake.date_time_between(start_date=start_datetime, end_date=end_datetime),
            'open_id': random.choice(open_id_list),
            'oid': random.choice(oid_list),
            'version': random.randint(0, 10)
        }
        # # 模拟一个用户多次与同一组件进行交互
        # insert_times = random.randint(1, 6)
        # for i in range(insert_times):
        data_list.append(post_data.copy())
        progress_bar.update(1)  # 更新进度条
    progress_bar.close() # 关闭进度条
    return data_list


def insert_batch_data(data_list, mongo_collection, batch_size=10000):
    '''
    将一个大的列表分批插入到数据库
    output: 插入的数据量大小
    '''
    print('开始插入数据：')
    total = len(data_list)
    progress_bar = tqdm(desc='插入数据', unit='条')
    for i in range(0, total, batch_size):
        batch = data_list[i:i+batch_size]
        try:
            mongo_collection.insert_many(batch)
            progress_bar.update(len(batch)) # 更新进度条
            print(f'Inserted batch {i//batch_size + 1}/{(total // batch_size) + 1}')
        except BulkWriteError as bwe:
            print(bwe.details)  # 打印详细的错误信息
            break  # 发生错误时跳出循环
        except Exception as e:
            print(f'General Error inserting batch {i//batch_size + 1}: {e}')
            break  # 发生错误时跳出循环
    progress_bar.close()  # 关闭进度条
    print('完成插入')
    return total



def transfer_collection_data(origin_collection, targeet_collection, data_size):
    '''
    从源数据库中随机获取指定数量的数据,并分批插入到数据库
    '''
    origin_records = origin_collection.aggregate([{'$sample': {'size': data_size}}])
    data_list = []
    progress_bar = tqdm(total=data_size, desc='取出数据', unit='条')
    for record in origin_records: # 迭代记录，并分批插入到数据库
        del record['_id'] # 删除原有记录的_id字段，MongoDB会自动生成新的_id字段
        data_list.append(record)
        progress_bar.update(1)  # 更新进度条
    progress_bar.close()  # 关闭进度条
    # 分批插入数据
    insert_batch_data(data_list=data_list, mongo_collection=targeet_collection, batch_size=100000)




if __name__ == "__main__":

    db = mongo_client['novapps_server']
    view_collection = db['ios_widget_analysis_widget_view']
    click_collection = db['ios_widget_analysis_widget_click']
    complete_set_collection = db['ios_widget_analysis_widget_complete_set']
    module_list_col = mongo_client['novapps_server']['ios_widget_component_library_module_list']

    results = module_list_col.find({},{'_id':1})

    fake = Faker()

    # oids = [ObjectId() for _ in range(OID_COUNT)] #oids: 包含了随机生成指定数量且为ObjectId类型的oid

    oids = [f['_id'] for f in results]
    print(len(oids))
    open_ids = [fake.uuid4() for _ in range(100000)] #open_ids: 包含了随机生成指定数量uuid作为open_id
    start_datetime = datetime(2023, 12, 5)
    end_datetime = datetime(2023, 12, 20)
    # 模拟展示量数据
    view_collection.drop() # 已清空展示量数据库
    print("展示量数据库开始插入")
    batch_data_toView = generate_bach_data(start_datetime=start_datetime, 
                                    end_datetime=end_datetime, 
                                    data_size=10000000,   # 生成数据量
                                    oid_list=oids,  # 组件列表
                                    open_id_list=open_ids,
                                    fake=fake) # 用户列表
    count = insert_batch_data(data_list=batch_data_toView, mongo_collection=view_collection, batch_size=100000) # 分批插入
    print(f"展示量数据库完成插入{count}条数据")

    ## 模拟点击量数据
    click_collection.drop() # 已清空点击量数据库
    print("点击量数据库开始插入")
    click_data_size = random.randint(8000000, count)
    transfer_collection_data(origin_collection=view_collection, targeet_collection=click_collection, data_size=click_data_size)
    print(f"点击量数据库完成插入{click_data_size}条数据")

    ## 模拟完成设置量数据
    complete_set_collection.drop() # 已清空完成设置量量数据库
    print("完成设置量数据库开始插入")
    set_data_size = random.randint(5000000, click_data_size)
    transfer_collection_data(origin_collection=click_collection, targeet_collection=complete_set_collection, data_size=set_data_size)
    # set_extra_data_size = random.randint(300, 400)
    # set_extra_data = generate_bach_data(start_datetime=start_datetime, 
    #                                 end_datetime=end_datetime, 
    #                                 data_size=set_extra_data_size,   # 生成数据量
    #                                 oid_list=oids,  # 组件列表
    #                                 open_id_list=open_ids,
    #                                 fake=fake) # 用户列表
    # count = insert_batch_data(data_list=set_extra_data, mongo_collection=complete_set_collection, batch_size=1000) # 分批插入
    # count = set_data_size + count
    print(f"完成设置量数据库完成插入{set_data_size}条数据")
